#pragma once
#include <unordered_map>
#include <mutex>
#include <stdexcept>

template <typename KeyA, typename KeyB>
class BiDirectionMap
{
private:
	std::unordered_map<KeyA, KeyB> forward;
	std::unordered_map<KeyB, KeyA> reverse;
	mutable std::mutex mapMutex;

public:
	size_t size() {
		std::lock_guard<std::mutex> lock(mapMutex);
		return forward.size();
	}

	void clear() {
		std::lock_guard<std::mutex> lock(mapMutex);
		forward.clear();
		reverse.clear();
	}

	void insert(const KeyA& a, const KeyB& b)
	{
		std::lock_guard<std::mutex> lock(mapMutex);

		if (forward.count(a) || reverse.count(b)) {
			throw std::runtime_error("Key already exists in bidirection map");
		}

		forward[a] = b;
		reverse[b] = a;
	}

	bool removeByA(const KeyA& a)
	{
		std::lock_guard<std::mutex> lock(mapMutex);

		auto itForward = forward.find(a);
		if (itForward == forward.end())
			return false;

		KeyB b = itForward->second;
		forward.erase(itForward);

		auto itReverse = reverse.find(b);
		if (itReverse != reverse.end())
			reverse.erase(itReverse);

		return true;
	}

	bool removeByB(const KeyB& b)
	{
		std::lock_guard<std::mutex> lock(mapMutex);

		auto itReverse = reverse.find(b);
		if (itReverse == reverse.end())
			return false;

		KeyA a = itReverse->second;
		reverse.erase(itReverse);

		auto itForward = forward.find(a);
		if (itForward != forward.end())
			forward.erase(itForward);

		return true;
	}

	KeyB getBFromA(const KeyA& a) const
	{
		std::lock_guard<std::mutex> lock(mapMutex);

		auto it = forward.find(a);
		if (it == forward.end())
			throw std::out_of_range("Key not found in forward map");

		return it->second;
	}

	KeyA getAFromB(const KeyB& b) const
	{
		std::lock_guard<std::mutex> lock(mapMutex);

		auto it = reverse.find(b);
		if (it == reverse.end())
			throw std::out_of_range("Key not found in reverse map");

		return it->second;
	}

	bool containsA(const KeyA& a) const
	{
		std::lock_guard<std::mutex> lock(mapMutex);
		return forward.count(a) > 0;
	}

	bool containsB(const KeyB& b) const
	{
		std::lock_guard<std::mutex> lock(mapMutex);
		return reverse.count(b) > 0;
	}
};

